#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : microbiome_beta_parser.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/10/15
@Description : 将 beta 多样性矩阵转换为长表格格式，并保留组内样本对
"""

from argparse import ArgumentParser, Namespace
from collections import defaultdict
from pathlib import Path

import pandas as pd


def _get_group_info(group_file: str) -> dict[str, str]:
    """
    读取分组信息文件，并将样本-组对应关系保存到字典中

    :param group_file: 分组文件路径，文件第一列为样本名，第二列为组名
    :return: 字典 {sample: group}
    """
    group_dict = {}
    with open(group_file, "r", encoding="utf-8") as f:
        lines = f.readlines()[1:]  # 跳过表头
        for line in lines:
            if line.strip():  # 忽略空行
                parts = line.strip().split()
                if len(parts) >= 2:
                    sample, group = parts[0], parts[1]
                    group_dict[sample] = group
    return group_dict


def _get_group_within_pairs(beta_df: pd.DataFrame, group_dict: dict[str, str]) -> pd.DataFrame:
    """
    将 beta 多样性矩阵整理成长表格格式，只保留组内样本对

    :param beta_df: beta 距离矩阵，行列都是样本名
    :param group_dict: 样本-组对应字典
    :return: DataFrame，包括列 ['X', 'Y', 'Group', 'Distance']
    """
    out_dict = defaultdict(list)
    seen_pairs = set()  # 用集合记录已处理的样本对，避免重复

    samples = beta_df.index
    for i in samples:
        for j in samples:
            # 样本对唯一性：使用 frozenset，保证 {i,j} 和 {j,i} 视为同一对
            pair = frozenset([i, j])
            if pair in seen_pairs:
                continue
            seen_pairs.add(pair)

            # 只保留组内样本对
            if group_dict.get(i) == group_dict.get(j):
                out_dict["X"].append(i)
                out_dict["Y"].append(j)
                out_dict["Group"].append(group_dict[i])
                out_dict["Distance"].append(beta_df.loc[i, j])

    return pd.DataFrame(out_dict)


def main(args: Namespace):
    """
    主函数：读取 beta 矩阵和分组文件，生成组内样本对长表格

    :param args: 命令行参数
    """
    output_path = Path(args.output_file).resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # 读取 beta 多样性矩阵
    beta_df = pd.read_csv(args.beta_matrix, sep="\t", index_col=0)

    # 读取分组信息
    group_dict = _get_group_info(args.group_table)

    # 生成组内样本对
    out_df = _get_group_within_pairs(beta_df, group_dict)

    # 输出结果
    out_df.to_csv(output_path, sep="\t", index=False)
    print(f"✅ 输出完成: {output_path}")


if __name__ == "__main__":
    parser = ArgumentParser(description="解析 beta 多样性矩阵为组内样本对长表格")
    parser.add_argument("--beta_matrix", type=str, required=True,
                        help="beta 多样性矩阵文件路径，行列都是样本名")
    parser.add_argument("--group_table", type=str, required=True,
                        help="样本分组文件路径，第一列样本名，第二列分组名")
    parser.add_argument("--output_file", type=str, required=True,
                        help="输出文件路径（长表格格式）")
    args = parser.parse_args()
    main(args)
